#include "modified_pairwise.hpp"

#include "MG.hpp"
#include "parcsr_matop.hpp"

template<typename idx_t, typename setup_t, int dof>
idx_t* modified_pairwise(const par_CSRMatrix<idx_t, setup_t, setup_t, dof> * par_fine_mat,
    par_CSRMatrix<idx_t, setup_t, setup_t, 1> * & R_mat, par_CSRMatrix<idx_t, setup_t, setup_t, 1> * & P_mat,
    const idx_t numRd)
{
    const MPI_Comm comm = par_fine_mat->comm;
    int my_pid; MPI_Comm_rank(comm, & my_pid);
    int num_procs; MPI_Comm_size(comm, & num_procs);
    constexpr MPI_Datatype mpi_idx_type = sizeof(idx_t) == 8 ? MPI_LONG_LONG : MPI_INT;

    const idx_t my_F_beg = par_fine_mat->beg_row;
    const idx_t my_num_F = par_fine_mat->end_row - par_fine_mat->beg_row;
    auto trans_loc2glb = [my_F_beg](const idx_t loc) {
        return loc + my_F_beg;
    };

    par_CSRMatrix<idx_t, setup_t, setup_t, 1> S_mat(comm, par_fine_mat->rows_partition, par_fine_mat->cols_partition);
    S_mat.commpkg = par_fine_mat->commpkg;
    S_mat.col_map_offd = par_fine_mat->col_map_offd;
    assert(S_mat.beg_row == S_mat.beg_col && S_mat.end_row == S_mat.end_col);

    constexpr int e_size = dof*dof;
    auto calc_blk_norm = [] (const setup_t * val) {
        setup_t ret = 0.0;
        for (int f = 0; f < e_size; f++) ret += val[f] * val[f];
        // ret = val[4]*val[4] + val[5]*val[5] + val[7]*val[7] + val[8]*val[8]; assert(dof == 3);
        return sqrt(ret);
    };
    std::vector<idx_t> bdr_F;// 细网格上位于进程边界的点，不含孤立井
    bool sorted = false;
    auto sort_and_find_bdr = [&S_mat, my_num_F, my_F_beg, & bdr_F, & sorted] () {
        for (idx_t i = 0; i < my_num_F; i++) {
#define PRES_COL_ORDER
#ifdef PRES_COL_ORDER // 保证列号升序（以便coupled和naive在单进程时严格复现）
            sort_vals_with_indices<idx_t, setup_t>(S_mat.diag.row_ptr[i + 1] - S_mat.diag.row_ptr[i],
                    S_mat.diag.vals    + S_mat.diag.row_ptr[i],
                    S_mat.diag.col_idx + S_mat.diag.row_ptr[i] );
#else // 快排不稳定排序
            my_qsort1<idx_t, setup_t, idx_t, -1>(   S_mat.diag.vals    + S_mat.diag.row_ptr[i],
                                                    S_mat.diag.col_idx + S_mat.diag.row_ptr[i],
                                        0, S_mat.diag.row_ptr[i + 1] - S_mat.diag.row_ptr[i] - 1);
#endif
            if (S_mat.offd.row_ptr[i + 1] > S_mat.offd.row_ptr[i]) {
#ifdef PRES_COL_ORDER // 保证列号升序（以便coupled和naive在单进程时严格复现）
                sort_vals_with_indices<idx_t, setup_t>(S_mat.offd.row_ptr[i + 1] - S_mat.offd.row_ptr[i],
                    S_mat.offd.vals    + S_mat.offd.row_ptr[i],
                    S_mat.offd.col_idx + S_mat.offd.row_ptr[i] );
#else // 快排不稳定排序
                my_qsort1<idx_t, setup_t, idx_t, -1>(   S_mat.offd.vals    + S_mat.offd.row_ptr[i],
                                                        S_mat.offd.col_idx + S_mat.offd.row_ptr[i],
                                            0, S_mat.offd.row_ptr[i + 1] - S_mat.offd.row_ptr[i] - 1);
#endif
                // 如果该点是井，则不算入边界点，因为它们不参与聚合，直接作为粗点
                if (iso_wells.find(i + my_F_beg) != iso_wells.end()) continue;
                else bdr_F.push_back(i);// 记录这个点为处于进程边界的点
            }
        }
        sorted = true;
    };


    S_mat.diag.nrows = par_fine_mat->diag.nrows; S_mat.offd.nrows = par_fine_mat->offd.nrows;
    S_mat.diag.ncols = par_fine_mat->diag.ncols; S_mat.offd.ncols = par_fine_mat->offd.ncols;
    S_mat.diag.nnz   = par_fine_mat->diag.nnz  ; S_mat.offd.nnz   = par_fine_mat->offd.nnz  ;
    S_mat.diag.alloc_mem(false);                 S_mat.offd.alloc_mem(false);
    for (idx_t i = 0; i <= my_num_F; i++) {
        S_mat.diag.row_ptr[i] = par_fine_mat->diag.row_ptr[i];
        S_mat.offd.row_ptr[i] = par_fine_mat->offd.row_ptr[i];
    }
    for (idx_t p = 0; p <  S_mat.diag.row_ptr[my_num_F]; p ++) {
        S_mat.diag.col_idx[p] = par_fine_mat->diag.col_idx[p];
        S_mat.diag.vals   [p] = calc_blk_norm(par_fine_mat->diag.vals + p * e_size);
    }
    for (idx_t p = 0; p <  S_mat.offd.row_ptr[my_num_F]; p ++) {
        S_mat.offd.col_idx[p] = par_fine_mat->offd.col_idx[p];
        S_mat.offd.vals   [p] = calc_blk_norm(par_fine_mat->offd.vals + p * e_size);
    }
    sort_and_find_bdr();

    assert(sorted);
    par_CSRMatrix<idx_t, setup_t, setup_t, 1> S_trans(comm);
    par_CSRMatrixTranspose(S_mat, S_trans, false);

    idx_t * loc_map_f2c = new idx_t [my_num_F];
    idx_t * offd_status = new idx_t [S_mat.offd.ncols];
    int   * offd_pids   = new int   [S_mat.offd.ncols];
    coar_wells.clear();// 为这一轮粗化清空
    constexpr idx_t UNAGGREGATED = -1, PENDING = -2, AGGREGATED_BY_OTHERS = -3;
    // 初始化
    for (idx_t fi = 0; fi < my_num_F; fi++)
        loc_map_f2c[fi] = UNAGGREGATED;
    for (idx_t j = 0; j < S_mat.offd.ncols; j ++)
        offd_status[j] = UNAGGREGATED;
    for (idx_t r = 0; r < S_mat.commpkg->num_recvs; r ++) {
        const int pid = S_mat.commpkg->recv_pids[r];
        for (idx_t j = S_mat.commpkg->recv_vec_starts[r]; j < S_mat.commpkg->recv_vec_starts[r + 1]; j ++)
            offd_pids[j] = pid;
    }
    std::unordered_map<idx_t, idx_t> col_map_glb_FT;
    for (idx_t lojT = 0; lojT < S_trans.offd.ncols; lojT ++) {// lojT: local offd j of S_trans
        const idx_t glb_fjT = S_trans.col_map_offd[lojT];
        col_map_glb_FT.emplace(glb_fjT, lojT);
    }

    idx_t my_num_C = 0;
    for (idx_t iter = 0; iter < numRd; iter ++) {// 反复进行若干轮迭代
        // 邻居通信更新 offd_status
        idx_t * S_cpkg_sbuf = new idx_t [S_mat.commpkg->send_map_starts[S_mat.commpkg->num_sends]];
        S_mat.commpkg->init_update_halo(1, loc_map_f2c, S_cpkg_sbuf, offd_status);
        S_mat.commpkg->wait_update_halo();
        delete [] S_cpkg_sbuf; S_cpkg_sbuf = nullptr;

        std::unordered_map<idx_t, idx_t> requests;
        for (idx_t loc_fi = 0; loc_fi < my_num_F; loc_fi ++) {// 在当前轮中遍历所有点一遍
            if (loc_map_f2c[loc_fi] != UNAGGREGATED) continue;// 之前已经聚合过了

            const idx_t glb_fi = trans_loc2glb(loc_fi);
            if (iso_wells.find(glb_fi) != iso_wells.end()) {// 该点是孤立点，自成粗点
                assert(loc_map_f2c[loc_fi] == UNAGGREGATED);// 之前不能被聚合过
                coar_wells.emplace(my_num_C);
                loc_map_f2c[loc_fi] = my_num_C ++;
                continue;
            }
            // 从大到小考察邻居
            // 找出diag内最强耦合的邻居
            idx_t diag_candi = -1; setup_t diag_v = 0.0;
            for (idx_t p = S_mat.diag.row_ptr[loc_fi]; p < S_mat.diag.row_ptr[loc_fi + 1]; p ++) {// 从大到小考察邻居
                const idx_t loc_fj = S_mat.diag.col_idx[p];
                const idx_t glb_fj = trans_loc2glb(loc_fj);
                if (loc_fj != loc_fi && // 该邻居不能是自己
                    iso_wells.find(glb_fj) == iso_wells.end() && // 也不能是井
                    loc_map_f2c[loc_fj] == UNAGGREGATED) // 也尚未被聚合
                {
                    diag_candi = glb_fj;
                    diag_v     = S_mat.diag.vals[p];
                    break;
                }
            }
            // 找出offd内最强耦合的邻居
            idx_t offd_candi = -1; setup_t offd_v = 0.0; idx_t od_ptr = -1;
            for (idx_t p = S_mat.offd.row_ptr[loc_fi]; p < S_mat.offd.row_ptr[loc_fi + 1]; p ++) {// 从大到小考察邻居
                const idx_t loc_fj = S_mat.offd.col_idx[p];
                const idx_t glb_fj = S_mat.col_map_offd[loc_fj];
                if (offd_status[loc_fj] == UNAGGREGATED && // 该邻居尚未被聚合
                    iso_wells.find(glb_fj) == iso_wells.end() )// 也不能是井
                {
                    offd_candi = glb_fj;
                    offd_v     = S_mat.offd.vals[p];
                    od_ptr     = loc_fj;
                    break;
                }
            }
            idx_t final_candi = -1;// 最终确定的聚合邻居
            if (diag_candi != -1 && offd_candi != -1) {
                if (diag_v >= offd_v) final_candi = diag_candi;// 选值大的
                else                  final_candi = offd_candi;
            }
            else { final_candi = MAX(diag_candi, offd_candi); }
            if (final_candi == -1) {// 找不到合适的邻居进行聚合
                // nothing to do
            } else {
                if (final_candi == diag_candi) {// 聚合的邻居是自己本进程内的
                    const idx_t loc_fj = diag_candi - my_F_beg;
                    loc_map_f2c[loc_fi] = my_num_C;
                    loc_map_f2c[loc_fj] = my_num_C ++;
                } else {// 聚合的邻居位于别的进程
                    const bool odd_round = (iter % 2) == 1;
                    if ((  odd_round  && glb_fi > offd_candi) ||// 第奇数次轮：允许向点号比glb_fi更小的点发起聚合请求
                        ((!odd_round) && glb_fi < offd_candi) ) // 第偶数次轮：允许向点号比glb_fi更大的点发起聚合请求
                    {
                        offd_status[od_ptr] = AGGREGATED_BY_OTHERS - glb_fi;// 本地“认为”该邻居已被聚合了
                        loc_map_f2c[loc_fi] = PENDING;
                        requests.emplace(loc_fi, AGGREGATED_BY_OTHERS - offd_candi);
                    }
                }
            }
        }
        // 发送请求
        idx_t * send_buf = new idx_t [S_trans.commpkg->send_map_starts[S_trans.commpkg->num_sends]];
        for (idx_t s = 0; s < S_trans.commpkg->num_sends; s ++) {
            idx_t       * sbuf   = send_buf                        + S_trans.commpkg->send_map_starts[s];
            const idx_t * my_ids = S_trans.commpkg->send_map_elmts + S_trans.commpkg->send_map_starts[s];
            const idx_t num = S_trans.commpkg->send_map_starts[s + 1] - S_trans.commpkg->send_map_starts[s];
            for (idx_t k = 0; k < num; k++) {
                idx_t loc_fi = my_ids[k];
                if (requests.find(loc_fi) != requests.end()) // 在这一轮中该点需要发出请求
                {
                    sbuf[k] = AGGREGATED_BY_OTHERS - requests[loc_fi];// 恢复为offd_candi
                } else {// 无聚合请求
                    sbuf[k] = -2;
                }
            }
            MPI_Isend(sbuf, num, mpi_idx_type, S_trans.commpkg->send_pids[s], 113 + iter * 2, comm, S_trans.commpkg->send_reqs + s);
        }
        // 接收、处理、并回复请求
        idx_t * recv_buf = new idx_t [S_trans.offd.ncols];
        for (idx_t r = 0; r < S_trans.commpkg->num_recvs; r++) {
            idx_t * rbuf = recv_buf + S_trans.commpkg->recv_vec_starts[r];
            const idx_t num = S_trans.commpkg->recv_vec_starts[r + 1] - S_trans.commpkg->recv_vec_starts[r];
            MPI_Recv(rbuf, num, mpi_idx_type, S_trans.commpkg->recv_pids[r], 113 + iter * 2, comm, MPI_STATUS_IGNORE);
            for (idx_t k = 0; k < num; k ++) {
                if (rbuf[k] == -2) continue;// 无聚合请求，回复也继续用-2即可
                idx_t glb_fj = S_trans.col_map_offd[k + S_trans.commpkg->recv_vec_starts[r]];
                idx_t glb_fi = rbuf[k];
                idx_t loc_fi = glb_fi - my_F_beg;
                if (0 <= loc_fi && loc_fi < my_num_F) {// 请求所希望聚合的点确实是自己负责的
                    if (loc_map_f2c[loc_fi] == UNAGGREGATED)// 且其未聚合
                    {
                        loc_map_f2c[loc_fi] = AGGREGATED_BY_OTHERS - glb_fj;// 记录被谁所聚合
                        rbuf[k] = 1;// 原位修改缓冲区
                    } else
                        rbuf[k] = -1;// 请求失败
                } else // 否则，仍然将回复维持为-2
                    rbuf[k] = -2;
            }
            // 原路发回
            MPI_Isend(rbuf, num, mpi_idx_type, S_trans.commpkg->recv_pids[r], 114 + iter * 2, comm, S_trans.commpkg->recv_reqs + r);
        }
        // 确认请求
        idx_t * ack_buf = new idx_t [S_trans.commpkg->send_map_starts[S_trans.commpkg->num_sends]];
        for (idx_t s = 0; s < S_trans.commpkg->num_sends; s ++) {
            idx_t       * abuf   = ack_buf                         + S_trans.commpkg->send_map_starts[s];
            const idx_t * sbuf   = send_buf                        + S_trans.commpkg->send_map_starts[s];
            const idx_t * my_ids = S_trans.commpkg->send_map_elmts + S_trans.commpkg->send_map_starts[s];
            const idx_t num = S_trans.commpkg->send_map_starts[s + 1] - S_trans.commpkg->send_map_starts[s];
            MPI_Recv(abuf, num, mpi_idx_type, S_trans.commpkg->send_pids[s], 114 + iter * 2, comm, MPI_STATUS_IGNORE);
            for (idx_t k = 0; k < num; k ++) {
                if (sbuf[k] == -2 || abuf[k] == -2) { continue; }// 之前无聚合请求，或者所请求的不是该进程
                idx_t loc_fi = my_ids[k];
                assert(requests[loc_fi] == AGGREGATED_BY_OTHERS - sbuf[k]);
                assert(loc_map_f2c[loc_fi] == PENDING);
                if (abuf[k] == 1)// 请求被允许了
                {
                    loc_map_f2c[loc_fi] = my_num_C ++;// 新聚合的粗点由我（即提出请求的进程）负责
                } else { assert(abuf[k] == -1);
                    loc_map_f2c[loc_fi] = UNAGGREGATED;// 重新将状态修改为未聚合
                }
            }
        }
        MPI_Waitall(S_trans.commpkg->num_sends, S_trans.commpkg->send_reqs, MPI_STATUSES_IGNORE);
        MPI_Waitall(S_trans.commpkg->num_recvs, S_trans.commpkg->recv_reqs, MPI_STATUSES_IGNORE);
        delete [] send_buf; delete [] recv_buf; delete [] ack_buf;
    }

    idx_t * C_rows_partits = new idx_t [num_procs + 1];
    C_rows_partits[0] = 0;
    MPI_Allgather(& my_num_C, 1, mpi_idx_type, C_rows_partits + 1, 1, mpi_idx_type, comm);
    for (int p = 0; p < num_procs; p++)
        C_rows_partits[p + 1] += C_rows_partits[p];

    // 构造限制矩阵
    R_mat = new par_CSRMatrix<idx_t, setup_t, setup_t, 1>(comm, C_rows_partits, par_fine_mat->rows_partition);
    P_mat = new par_CSRMatrix<idx_t, setup_t, setup_t, 1>(comm, par_fine_mat->cols_partition, C_rows_partits);

    const idx_t _cbeg = C_rows_partits[my_pid], _cend = C_rows_partits[my_pid + 1];
    idx_t * glb_map_f2c = new idx_t [my_num_F];// loc_f => glb_c
    for (idx_t loc_fi = 0; loc_fi < my_num_F; loc_fi ++) {
        if (loc_map_f2c[loc_fi] >= 0)
            glb_map_f2c[loc_fi] = loc_map_f2c[loc_fi] + _cbeg;// 累加全局偏移
        else
            glb_map_f2c[loc_fi] = -1;// 暂时置为-1
    }
    {// 告知邻居进程，细网格上我的边界点的全局序号
        idx_t * send_buf = new idx_t [S_trans.commpkg->send_map_starts[S_trans.commpkg->num_sends]];
        idx_t * recv_buf = new idx_t [S_trans.offd.ncols];
        for (idx_t s = 0; s < S_trans.commpkg->num_sends; s++) {
            idx_t       * buf    = send_buf                        + S_trans.commpkg->send_map_starts[s];
            const idx_t * my_ids = S_trans.commpkg->send_map_elmts + S_trans.commpkg->send_map_starts[s];
            const idx_t num = S_trans.commpkg->send_map_starts[s + 1] - S_trans.commpkg->send_map_starts[s];
            for (idx_t k = 0; k < num; k++) {
                buf[k] = glb_map_f2c[my_ids[k]];
            }
            MPI_Isend(buf, num, mpi_idx_type, S_trans.commpkg->send_pids[s], 115 + numRd * 2, comm, S_trans.commpkg->send_reqs + s);
        }
        for (idx_t r = 0; r < S_trans.commpkg->num_recvs; r++) {
            idx_t * buf = recv_buf + S_trans.commpkg->recv_vec_starts[r];
            const idx_t num = S_trans.commpkg->recv_vec_starts[r + 1] - S_trans.commpkg->recv_vec_starts[r];
            MPI_Irecv(buf, num, mpi_idx_type, S_trans.commpkg->recv_pids[r], 115 + numRd * 2, comm, S_trans.commpkg->recv_reqs + r);
        }
        MPI_Waitall(S_trans.commpkg->num_recvs, S_trans.commpkg->recv_reqs, MPI_STATUSES_IGNORE);

        for (idx_t loc_fi = 0; loc_fi < my_num_F; loc_fi++) {
            if (loc_map_f2c[loc_fi] <= AGGREGATED_BY_OTHERS) { assert(glb_map_f2c[loc_fi] == -1);
                const idx_t glb_fj = - (loc_map_f2c[loc_fi] - AGGREGATED_BY_OTHERS);
                assert(col_map_glb_FT.find(glb_fj) != col_map_glb_FT.end());
                const idx_t loc_fj = col_map_glb_FT[glb_fj];
                const idx_t glb_cj = recv_buf[ loc_fj ];
                glb_map_f2c[loc_fi] = glb_cj;// 更新全局映射
            }
            // 此时先不管 loc_map_f2c[loc_fi] == UNAGGREGATED 即尚未聚合的点
        }
        MPI_Waitall(S_trans.commpkg->num_sends, S_trans.commpkg->send_reqs, MPI_STATUSES_IGNORE);
        delete [] recv_buf; delete [] send_buf;
    }
    {// 处理尚未聚合的点
        idx_t * send_buf = new idx_t [par_fine_mat->commpkg->send_map_starts[par_fine_mat->commpkg->num_sends]];
        idx_t * recv_buf = new idx_t [par_fine_mat->offd.ncols];
        for (idx_t s = 0; s < par_fine_mat->commpkg->num_sends; s++) {
            idx_t       * buf    = send_buf                              + par_fine_mat->commpkg->send_map_starts[s];
            const idx_t * my_ids = par_fine_mat->commpkg->send_map_elmts + par_fine_mat->commpkg->send_map_starts[s];
            const idx_t num = par_fine_mat->commpkg->send_map_starts[s + 1] - par_fine_mat->commpkg->send_map_starts[s];
            for (idx_t k = 0; k < num; k++) {
                buf[k] = glb_map_f2c[my_ids[k]];
            }
            MPI_Isend(buf, num, mpi_idx_type, par_fine_mat->commpkg->send_pids[s], 114, comm, par_fine_mat->commpkg->send_reqs + s);
        }
        for (idx_t r = 0; r < par_fine_mat->commpkg->num_recvs; r++) {
            idx_t * buf = recv_buf + par_fine_mat->commpkg->recv_vec_starts[r];
            const idx_t num = par_fine_mat->commpkg->recv_vec_starts[r + 1] - par_fine_mat->commpkg->recv_vec_starts[r];
            MPI_Irecv(buf, num, mpi_idx_type, par_fine_mat->commpkg->recv_pids[r], 114, comm, par_fine_mat->commpkg->recv_reqs + r);
        }
        MPI_Waitall(par_fine_mat->commpkg->num_recvs, par_fine_mat->commpkg->recv_reqs, MPI_STATUSES_IGNORE);
    
        for (idx_t loc_fi = 0; loc_fi < my_num_F; loc_fi++) {
            if (loc_map_f2c[loc_fi] == UNAGGREGATED) {// 该点仍未被聚合
                // 从所有邻居中选最大的进行聚合：注意包含落在别的进程上的邻居
                // 找出diag内最强耦合的邻居
                idx_t diag_candi = -1; setup_t diag_v = 0.0;
                for (idx_t p = S_mat.diag.row_ptr[loc_fi]; p < S_mat.diag.row_ptr[loc_fi + 1]; p ++) {// 从大到小考察邻居
                    const idx_t loc_fj = S_mat.diag.col_idx[p];
                    const idx_t glb_fj = trans_loc2glb(loc_fj);
                    if (loc_fj != loc_fi && // 该邻居不能是自己
                        iso_wells.find(glb_fj) == iso_wells.end() && // 也不能是井
                        glb_map_f2c[loc_fj] != -1)// 也不能未聚合
                    {
                        diag_candi = glb_map_f2c[loc_fj];
                        diag_v     = S_mat.diag.vals[p];
                        break;
                    }
                }
                // 找出offd内最强耦合的邻居
                idx_t offd_candi = -1; setup_t offd_v = 0.0;
                for (idx_t p = S_mat.offd.row_ptr[loc_fi]; p < S_mat.offd.row_ptr[loc_fi + 1]; p ++) {// 从大到小考察邻居
                    const idx_t loc_fj = S_mat.offd.col_idx[p];
                    const idx_t glb_fj = S_mat.col_map_offd[loc_fj];
                    if (iso_wells.find(glb_fj) == iso_wells.end() && // 不能是井
                        recv_buf[loc_fj] != -1) // 且已经被聚合了
                    {
                        offd_candi = recv_buf[loc_fj];
                        offd_v     = S_mat.offd.vals[p];
                        break;
                    }
                }
                idx_t final_candi = -1;// 最终确定的聚合邻居
                if (diag_candi != -1 && offd_candi != -1) {
                    if (diag_v > offd_v) final_candi = diag_candi;// 选值大的
                    else                 final_candi = offd_candi;
                }
                else { final_candi = MAX(diag_candi, offd_candi); }
                // if (loc_fi + _fbeg == 405) printf("Proc %d glb_i %d diag_v %.2f offd_v %.2f\n", my_pid, loc_fi + _fbeg, diag_v, offd_v);
                glb_map_f2c[loc_fi] = final_candi;
                // if (final_candi == -1) {// check: 此时还尚未被聚合的，一定是只有一个邻居（即自己的）
                //     assert(S_mat.diag.row_ptr[loc_fi + 1] - S_mat.diag.row_ptr[loc_fi] == 1);
                //     assert(S_mat.diag.col_idx[S_mat.diag.row_ptr[loc_fi]] == loc_fi);
                // }
            }
        }
        MPI_Waitall(par_fine_mat->commpkg->num_sends, par_fine_mat->commpkg->send_reqs, MPI_STATUSES_IGNORE);
        delete [] recv_buf; delete [] send_buf;
    }// 此时在glb_map_f2c[]中仍为-1的点将不被粗化进入下一层粗网格

    {// 构造插值矩阵
        P_mat->diag.nrows = my_num_F; P_mat->offd.nrows = my_num_F;
        P_mat->diag.row_ptr = new idx_t [my_num_F + 1]; P_mat->diag.row_ptr[0] = 0;
        P_mat->offd.row_ptr = new idx_t [my_num_F + 1]; P_mat->offd.row_ptr[0] = 0;
        std::unordered_map<idx_t, std::pair<int, idx_t> > Poffd_g2l;// 记录插值矩阵的非对角部分的全局序到局部序的映射
        for (idx_t loc_fi = 0; loc_fi < my_num_F; loc_fi ++) {
            const idx_t glb_ci = glb_map_f2c[loc_fi];
            if (_cbeg <= glb_ci && glb_ci < _cend) {// 粗点属于本进程负责计算
                P_mat->diag.row_ptr[loc_fi + 1] = 1;
                P_mat->offd.row_ptr[loc_fi + 1] = 0;
            } else if (glb_ci >= 0) {// 粗点不属于本进程计算
                P_mat->diag.row_ptr[loc_fi + 1] = 0;
                P_mat->offd.row_ptr[loc_fi + 1] = 1;
                
                std::pair<int, idx_t> tmp(-1, -1);
                Poffd_g2l.emplace(glb_ci, tmp);
            } else { assert(glb_ci == -1);// 粗点不进入下一层
                P_mat->diag.row_ptr[loc_fi + 1] = 0;
                P_mat->offd.row_ptr[loc_fi + 1] = 0;
            }
        }
        for (idx_t loc_fi = 0; loc_fi < my_num_F; loc_fi ++) {// prefix sum
            P_mat->diag.row_ptr[loc_fi + 1] += P_mat->diag.row_ptr[loc_fi];
            P_mat->offd.row_ptr[loc_fi + 1] += P_mat->offd.row_ptr[loc_fi];
        }
        P_mat->prepare_commpkg(Poffd_g2l);

        // 确定列数和非零元数
        P_mat->diag.ncols = my_num_C;                      P_mat->offd.ncols = Poffd_g2l.size();
        P_mat->diag.nnz   = P_mat->diag.row_ptr[my_num_F]; P_mat->offd.nnz   = P_mat->offd.row_ptr[my_num_F];
        // 分配列序号和值的内存
        P_mat->diag.col_idx = new idx_t  [P_mat->diag.nnz]; P_mat->offd.col_idx = new idx_t  [P_mat->offd.nnz];
        P_mat->diag.vals    = new setup_t[P_mat->diag.nnz]; P_mat->offd.vals    = new setup_t[P_mat->offd.nnz];

        for (idx_t loc_fi = 0; loc_fi < my_num_F; loc_fi ++) {
            const idx_t glb_ci = glb_map_f2c[loc_fi];
            // const idx_t dcnt = 0, ocnt = 0;
            const idx_t doff = P_mat->diag.row_ptr[loc_fi],
                        ooff = P_mat->offd.row_ptr[loc_fi];
            if (_cbeg <= glb_ci && glb_ci < _cend) {// 粗点属于本进程负责计算
                P_mat->diag.col_idx[doff] = glb_ci - _cbeg;
                P_mat->diag.vals   [doff] = 1.0;
            } else if (glb_ci >= 0) {// 粗点不属于本进程计算
                assert(Poffd_g2l.find(glb_ci) != Poffd_g2l.end());
                P_mat->offd.col_idx[ooff] = Poffd_g2l[glb_ci].second;
                P_mat->offd.vals   [ooff] = 1.0;
            } else { assert(glb_ci == -1);// 粗点不进入下一层
                // nothing to do
            }
        }
#ifdef DEBUG
        for (int p = 0; p < num_procs; p++) {
            if (my_pid == p) {
                printf("Proc %d Prlg diag %d %d %d offd %d %d %d\n", my_pid, P_mat->diag.nrows, P_mat->diag.ncols, P_mat->diag.nnz,
                    P_mat->offd.nrows, P_mat->offd.ncols, P_mat->offd.nnz);
                for (idx_t i = 0; i < P_mat->diag.nrows; i++) {
                    printf("  %d : ", P_mat->beg_row + i);
                    for (idx_t p = P_mat->diag.row_ptr[i]; p < P_mat->diag.row_ptr[i + 1]; p++)
                        printf("%d ", P_mat->beg_col + P_mat->diag.col_idx[p]);
                    printf(" | ");
                    for (idx_t p = P_mat->offd.row_ptr[i]; p < P_mat->offd.row_ptr[i + 1]; p++)
                        printf("%d ", P_mat->col_map_offd[P_mat->offd.col_idx[p]]);
                    printf("\n");
                }
            }
            MPI_Barrier(comm);
        }
#endif
    }
    par_CSRMatrixTranspose(* P_mat, * R_mat);
    {// 需要修改限制矩阵的值
        for (idx_t i = 0; i < R_mat->diag.nrows; i ++) {
            const idx_t num = R_mat->diag.row_ptr[i + 1] - R_mat->diag.row_ptr[i]
                            + R_mat->offd.row_ptr[i + 1] - R_mat->offd.row_ptr[i];
            const setup_t v = 1.0 / (setup_t) num;
            for (idx_t p = R_mat->diag.row_ptr[i]; p < R_mat->diag.row_ptr[i + 1]; p ++) {
                assert(R_mat->diag.vals[p] == 1.0);
                R_mat->diag.vals[p] = v;
            }
            for (idx_t p = R_mat->offd.row_ptr[i]; p < R_mat->offd.row_ptr[i + 1]; p ++) {
                assert(R_mat->offd.vals[p] == 1.0);
                R_mat->offd.vals[p] = v;
            }
        }
#ifdef DEBUG
        for (int p = 0; p < num_procs; p++) {
            if (my_pid == p) {
                printf("Proc %d Rstr diag %d %d %d offd %d %d %d\n", my_pid, R_mat->diag.nrows, R_mat->diag.ncols, R_mat->diag.nnz,
                    R_mat->offd.nrows, R_mat->offd.ncols, R_mat->offd.nnz);
                for (idx_t i = 0; i < R_mat->diag.nrows; i++) {
                    printf("  %d : ", R_mat->beg_row + i);
                    for (idx_t p = R_mat->diag.row_ptr[i]; p < R_mat->diag.row_ptr[i + 1]; p++)
                        printf("%d ", R_mat->beg_col + R_mat->diag.col_idx[p]);
                    printf(" | ");
                    for (idx_t p = R_mat->offd.row_ptr[i]; p < R_mat->offd.row_ptr[i + 1]; p++)
                        printf("%d ", R_mat->col_map_offd[R_mat->offd.col_idx[p]]);
                    printf("\n");
                }
            }
            MPI_Barrier(comm);
        }
#endif
    }

// #ifdef DEBUG
    const idx_t my_num_I = my_num_F - P_mat->diag.nnz - P_mat->offd.nnz;// 独立点个数
    idx_t tot_num_I = 0;
    MPI_Reduce(& my_num_I, & tot_num_I, 1, mpi_idx_type, MPI_SUM, 0, comm);
    if (my_pid == 0) printf("Glb #I %d\n", tot_num_I);
// #endif

    delete [] loc_map_f2c; delete [] glb_map_f2c;

    return C_rows_partits;
}

template int* modified_pairwise(const par_CSRMatrix_I32F64F64_DOF1 *, par_CSRMatrix_I32F64F64_DOF1 * &, par_CSRMatrix_I32F64F64_DOF1 * &, const int);
template int* modified_pairwise(const par_CSRMatrix_I32F64F64_DOF2 *, par_CSRMatrix_I32F64F64_DOF1 * &, par_CSRMatrix_I32F64F64_DOF1 * &, const int);
template int* modified_pairwise(const par_CSRMatrix_I32F64F64_DOF3 *, par_CSRMatrix_I32F64F64_DOF1 * &, par_CSRMatrix_I32F64F64_DOF1 * &, const int);
template int* modified_pairwise(const par_CSRMatrix_I32F64F64_DOF4 *, par_CSRMatrix_I32F64F64_DOF1 * &, par_CSRMatrix_I32F64F64_DOF1 * &, const int);

template int* modified_pairwise(const par_CSRMatrix_I32F32F32_DOF1 *, par_CSRMatrix_I32F32F32_DOF1 * &, par_CSRMatrix_I32F32F32_DOF1 * &, const int);
template int* modified_pairwise(const par_CSRMatrix_I32F32F32_DOF2 *, par_CSRMatrix_I32F32F32_DOF1 * &, par_CSRMatrix_I32F32F32_DOF1 * &, const int);
template int* modified_pairwise(const par_CSRMatrix_I32F32F32_DOF3 *, par_CSRMatrix_I32F32F32_DOF1 * &, par_CSRMatrix_I32F32F32_DOF1 * &, const int);
template int* modified_pairwise(const par_CSRMatrix_I32F32F32_DOF4 *, par_CSRMatrix_I32F32F32_DOF1 * &, par_CSRMatrix_I32F32F32_DOF1 * &, const int);

